#include "kernel_operator.h"

constexpr int32_t TOTAL_LENGTH = 8 * 10 * 10;                         // total length of data
constexpr int32_t USE_CORE_NUM = 8;                                   // num of core used
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM;         // length computed of each core
constexpr int32_t TILE_NUM = 8;                                       // split data into 8 tiles for each core
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; // separate to 2 parts, due to double buffer

class KernelTriuCustom
{
public:
    __aicore__ inline KernelTriuCustom() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int32_t diagVal)
    {
        this->diagVal = diagVal;
        xGm.SetGlobalBuffer((__gm__ half *)x, BLOCK_LENGTH);
        yGm.SetGlobalBuffer((__gm__ half *)y, BLOCK_LENGTH);
        pipe.InitBuffer(queBind, BUFFER_NUM, TILE_LENGTH * sizeof(half));
    }
    __aicore__ inline void Process()
    {
        int32_t cnt = 0;
        for (int32_t i = 0; i < 8; i++)
        {
            for (int32_t j = 0; j < 10; j++)
            {
                int32_t k = 0;
                while (k < 10 && k - j < this->diagVal)
                {
                    yGm.SetValue(cnt, (half)0);
                    k++;
                    cnt++;
                }
                while (k < 10)
                {
                    half curr = xGm.GetValue(cnt);
                    yGm.SetValue(cnt, curr);
                    k++;
                    cnt++;
                }
            }
        }
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQueBind<AscendC::QuePosition::VECIN, AscendC::QuePosition::VECOUT, BUFFER_NUM> queBind;
    AscendC::GlobalTensor<half> xGm;
    AscendC::GlobalTensor<half> yGm;

    int32_t diagVal = 0;
};

extern "C" __global__ __aicore__ void triu_custom(GM_ADDR x, GM_ADDR y)
{
    KernelTriuCustom op;
    op.Init(x, y, 0);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void triu_custom_do(uint32_t blockDim, void *stream, uint8_t *x, uint8_t *y)
{
    triu_custom<<<blockDim, nullptr, stream>>>(x, y);
}
#endif